import torch
from torch.utils.tensorboard import SummaryWriter
import torch.nn as nn
import torch.optim as optim
from models.VGG.VGGMain import VGG19
from utils import get_data_loaders

def train_model():
    train_loader, _ = get_data_loaders()  # 获取训练数据加载器

    model = VGG19().cuda()
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-5)

    # 训练
    log_dir = './train-log-01'
    writer = SummaryWriter(log_dir)
    best_loss = float('inf')

    for epoch in range(10):
        print(f"~第{epoch + 1}轮训练开始~")
        total_loss = 0
        for data in train_loader:
            imgs, labels = data
            imgs, labels = imgs.cuda(), labels.cuda()

            outputs = model(imgs)
            loss = loss_fn(outputs, labels)
            total_loss += loss.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        average_loss = total_loss / len(train_loader)
        writer.add_scalar("ave_loss/every_train_epoch", average_loss, epoch)

        if average_loss < best_loss:
            best_loss = average_loss
            torch.save(model.state_dict(), 'best_model_AlexNet.pth')

        print(f"~第{epoch + 1}轮训练结束~，平均损失函数为{average_loss}")

    writer.close()